/*
Package middlewares
@Author   : Uyynot
@Email    : uyynot@qq.com
@Time     : 2024/3/29 18:44
@File     : limit.go
@Project  : gohub
@Desc     :
*/
package middlewares

import (
	"github.com/gin-gonic/gin"
	"github.com/spf13/cast"
	"gohub/pkg/app"
	"gohub/pkg/limiter"
	"gohub/pkg/logger"
	"gohub/pkg/response"
	"net/http"
)

func LimitIP(limit string) gin.HandlerFunc {
	if app.IsTesting() {
		limit = "1000000-H"
	}
	return func(c *gin.Context) {
		//针对IP限流
		key := limiter.GetKeyIP(c)
		if ok := LimitHandler(c, key, limit); !ok {
			return
		}
		c.Next()
	}
}

func LimitPerRoute(limit string) gin.HandlerFunc {
	if app.IsTesting() {
		limit = "1000000-H"
	}
	return func(c *gin.Context) {
		//针对单个路由，增加访问次数
		c.Set("limiter-once", false)

		//针对IP+路由进行限流
		key := limiter.GetKeyRouteWithIP(c)
		if ok := LimitHandler(c, key, limit); !ok {
			return
		}
		c.Next()
	}
}

func LimitHandler(c *gin.Context, key, limit string) bool {
	//获取超额情况
	rate, err := limiter.CheckRate(c, key, limit)
	if err != nil {
		logger.LogIf(err)
		response.Abort500(c)
		return false
	}

	//设置标头信息
	//X-ReteLimit-Limit: 10000 最大访问次数
	//X-ReteLimit-Remaining: 9999 当前剩余访问次数
	//X-ReteLimit-Reset: 1616959474 到改时间点，访问次数会重置为X-RateLimit-Limit（即最大访问次数）
	c.Header("X-ReteLimit-Limit", cast.ToString(rate.Limit))
	c.Header("X-ReteLimit-Remaining", cast.ToString(rate.Remaining))
	c.Header("X-ReteLimit-Reset", cast.ToString(rate.Reset))

	//超额
	if rate.Reached {
		//"提示用户超额"
		c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{
			"message": "访问过于频繁，请稍后再试",
		})
		return false
	}
	return true
}
